#include "EncryptedStream.h"
#include "Config.h"

#include <stdio.h>
#include <string.h>

#include <sha.h>

EncryptedStream::EncryptedStream()
	: sock(NULL)
	, sockset(NULL)
	, buffer(4096)
{
}

EncryptedStream::~EncryptedStream() {
	destroy();
}

void EncryptedStream::destroy() {
	if (sockset) {
		SDLNet_TCP_DelSocket(sockset, sock);
		SDLNet_FreeSocketSet(sockset);
		sockset = NULL;
	}
	if (sock) {
		SDLNet_TCP_Close(sock);
		sock = NULL;
	}
}

void EncryptedStream::create(TCPsocket sock_) {
	if (sock_ && sock != sock_) {
		destroy();
		sock = sock_;
		sockset = SDLNet_AllocSocketSet(1);
		SDLNet_TCP_AddSocket(sockset, sock);

		enc = CryptoPP::XChaCha20::Encryption();
		dec = CryptoPP::XChaCha20::Decryption();

		std::vector<CryptoPP::byte> key, iv;

		{
			std::string s;
			CryptoPP::SHA256 hash;
			s = Config::getConfig("key");
			hash.Update((const CryptoPP::byte*)s.c_str(), s.size());
			key.resize(hash.DigestSize());
			hash.Final(key.data());
		}

		{
			std::string s;
			CryptoPP::SHA256 hash;
			s = Config::getConfig("iv");
			hash.Update((const CryptoPP::byte*)s.c_str(), s.size());
			iv.resize(hash.DigestSize());
			hash.Final(iv.data());
		}

		enc.SetKeyWithIV(key.data(), std::min(key.size(), enc.MaxKeyLength()), iv.data(), std::min(iv.size(), enc.MaxIVLength()));
		dec.SetKeyWithIV(key.data(), std::min(key.size(), enc.MaxKeyLength()), iv.data(), std::min(iv.size(), enc.MaxIVLength()));
	}
}

int EncryptedStream::send(const void *data, int len) {
	if (sock == NULL) return 0;

	int ret = 0;

	while (len > 0) {
		int i = std::min<int>(len, buffer.size());
		enc.ProcessData((CryptoPP::byte*)buffer.data(), (const CryptoPP::byte*)data, i);
		int j = SDLNet_TCP_Send(sock, buffer.data(), i);
		if (j > 0) {
			ret += j;
			len -= j;
			data = ((const char*)data) + j;
		}
		if (j < i) {
			printf("ERROR: failed to send data; GFW'ed\n");
			destroy();
			break;
		}
	}

	return ret;
}

int EncryptedStream::receive(void *data, int len, bool exact) {
	if (sock == NULL) return 0;

	int ret = 0;

	while (len > 0) {
		int i = std::min<int>(len, buffer.size());
		int j = SDLNet_TCP_Recv(sock, buffer.data(), i);
		if (j <= 0) {
			printf("ERROR: failed to receive data; GFW'ed\n");
			destroy();
			break;
		}
		dec.ProcessData((CryptoPP::byte*)data, (const CryptoPP::byte*)buffer.data(), j);
		ret += j;
		len -= j;
		data = ((char*)data) + j;

		if (!exact) break;
	}

	return ret;
}

void EncryptedStream::dropIncoming() {
	while (check(0)) {
		int j = SDLNet_TCP_Recv(sock, buffer.data(), buffer.size());
		if (j <= 0) {
			printf("ERROR: failed to receive data; GFW'ed\n");
			destroy();
			break;
		}
		dec.DiscardBytes(j);
	}
}

bool EncryptedStream::check(int timeout) {
	return sock && sockset && SDLNet_CheckSockets(sockset, timeout) > 0 && SDLNet_SocketReady(sock);
}
